#include "ppe/block_manager.h"
#include "common/common_types.h"
#include "NanoFlowVis/Substance.h"
#include "common/malloc_align.h"
#include "NanoFlowVis/myIO.h"

BlockManager::BlockManager(uint32_t kCount, uint32_t jCount, uint32_t iCount, uint32_t speCount)
: _kCount(kCount)
, _jCount(jCount)
, _iCount(iCount)
, _speCount(speCount)
, _blocks(NULL)
, _dummyInput(NULL)
{
    assert(speCount >= 1);

    _blocks = new BlockInfo**[_kCount + 1];
    for (int k = 0; k < _kCount + 1; ++k) {
        _blocks[k] = new BlockInfo*[_jCount + 1];
        for (int j = 0; j < _jCount + 1; ++j) {
            _blocks[k][j] = new BlockInfo[_iCount + 1];
            for (int i = 0; i < _iCount + 1; ++i) {
                _blocks[k][j][i].inputData  = (IncomingDataPacket*)malloc_align(sizeof(IncomingDataPacket), 7);
                for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
                    _blocks[k][j][i].outputData[shiftIdx] = (OutgoingDataPacket*)malloc_align(sizeof(OutgoingDataPacket), 7);
                }
            }
        }
    }
    _speInitDatas = (SpeInitData*)malloc_align(sizeof(SpeInitData)*_speCount, 7);
    _forces = (ForcesArray*)malloc_align(sizeof(ForcesArray), 7);
    _dummyInput = (IncomingDataPacket*)malloc_align(sizeof(IncomingDataPacket), 7);
}

void BlockManager::packData(Cell*** cells, float tau, IFr*** forces, int forcesCount) {
    IntCoords prevCoords[8];
    bool hasPrev[8];
    _ZeroStruct(&hasPrev);
    _ZeroStruct(&prevCoords);
    //printf("packing %d x %d x %d cube, arguments (0x%x, %f, 0x%x)\n", _kCount + 1, _jCount + 1, _iCount + 1, (int)(long long)cells, tau, (int)(long long)forces);
    for (int i = 0; i < forcesCount; ++i) {
        for (int j = 0; j < forcesCount; ++j) {
            if (forces[i][j]) {
//                 printf("  assigning to 0x%x[%d][%d] from 0x%x\n", (int)(long long)_forces, i, j, (int)(long long)forces[i][j]);
                _forces->array[i][j] = forces[i][j]->getForceInfo();
            } else {
//                 printf("  forces[%d][%d] if null\n", i, j);
            }
        }
    }
    for (int s = 0; s < _speCount; ++s) {
        _ZeroStruct(&_speInitDatas[s]);
        _speInitDatas[s].tau = tau;
        _speInitDatas[s].forcesInfoPtr = (unsigned long long)_forces;
    }
    int speNum = 0;
    for (int k = 0; k < _kCount + 1; ++k) {
        for (int j = 0; j < _jCount + 1; ++j) {
            for (int i = 0; i < _iCount + 1; ++i) {
//                 printf("  fill iteration; k = %d, j = %d, i = %d\n", k, j, i);
                int particlesInBatch = 0;
                for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
                    if (_isValidShift(k, j, i, shiftIdx)) {
                        particlesInBatch += SHIFT_IDX(cells, k, j, i, shiftIdx).curIn;
                    }
                }
                if (particlesInBatch == 0) {
//                     printf("    no particles in batch, continue...\n");
                    continue;
                } else {
//                     printf("    %d particles in batch\n", particlesInBatch);
                }
                IncomingDataHeader& header = _blocks[k][j][i].inputData->header;
                if (hasPrev[speNum]) {
                    const IntCoords& prev = prevCoords[speNum];
//                     printf("    prevCoords[%d] == {k:%d, j:%d, i:%d}\n", speNum, prev.k, prev.j, prev.i);
                    for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
                        header.nextIncomingDataPtr[shiftIdx] = (unsigned long long)_getInput(prev.k, prev.j, prev.i, shiftIdx);
                        header.nextOutgoingDataPtr[shiftIdx] = (unsigned long long)_blocks[prev.k][prev.j][prev.i].outputData[shiftIdx];
                    }
                } else {
                    _ZeroStruct(&header.nextIncomingDataPtr);
                    _ZeroStruct(&header.nextOutgoingDataPtr);
                    hasPrev[speNum] = true;
                }

//                 printf("    filling spe init data\n");

                for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
                    _speInitDatas[speNum].incomingDataPtr[shiftIdx] = (unsigned long long)_getInput(k, j, i, shiftIdx);
                    _speInitDatas[speNum].outgoingDataPtr[shiftIdx] = (unsigned long long)_blocks[k][j][i].outputData[shiftIdx];
                }

                ++_speInitDatas[speNum].packetCount;

                for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
                    if (!_isValidShift(k, j, i, shiftIdx)) {
                        continue;
                    }
                    Cell& cell = SHIFT_IDX(cells, k, j, i, shiftIdx);
                    IncomingDataPacket* inputData = SHIFT_IDX(_blocks, k, j, i, shiftIdx).inputData;

                    inputData->header.particlesCount = cell.curIn;
                    ParticleData* data = inputData->data;
                    for (int c = 0; c < cell.curIn; ++c) {
                        data[c].pos = cell.inCell[c]->getPos();
                        data[c].mass = cell.inCell[c]->getM();
                        data[c].componentId = cell.inCell[c]->getCompID();
                        data[c].moleculaId = cell.inCell[c]->getMolID();
                    }
                }

//                 printf("    replacing prevCoords[%d] == {k:%d, j:%d, i:%d} with {k:%d, j:%d, i:%d}\n", speNum, prevCoords[speNum].k, prevCoords[speNum].j, prevCoords[speNum].i, k, j, i);
                prevCoords[speNum].k = k;
                prevCoords[speNum].j = j;
                prevCoords[speNum].i = i;
                speNum = (speNum + 1) % _speCount;
            }
        }
    }
}

void BlockManager::addAccels(Cell*** cells) {
    for (int k = 0; k < _kCount + 1; ++k) {
        for (int j = 0; j < _jCount + 1; ++j) {
            for (int i = 0; i < _iCount + 1; ++i) {
                for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
                    if (!_isValidShift(k, j, i, shiftIdx)) {
                        continue;
                    }
                    Cell& cell = SHIFT_IDX(cells, k, j, i, shiftIdx);
                    const Vector3* accels = _blocks[k][j][i].outputData[shiftIdx]->accels;
                    for (int c = 0; c < cell.curIn; ++c) {
                        cell.inCell[c]->addAccel(accels[c]);
                    }
                }
            }
        }
    }
}

BlockManager::~BlockManager() {
    free_align(_dummyInput);
    free_align(_forces);
    free_align(_speInitDatas);
    for (int k = 0; k < _kCount + 1; ++k) {
        for (int j = 0; j < _jCount + 1; ++j) {
            for (int i = 0; i < _iCount + 1; ++i) {
                free_align(_blocks[k][j][i].inputData);
                for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
                    free_align(_blocks[k][j][i].outputData[shiftIdx]);
                }
            }
            delete[] _blocks[k][j];
        }
        delete[] _blocks[k];
    }
    delete[] _blocks;
}
